package com.tianji.aigc.service.impl;

import com.baomidou.mybatisplus.extension.service.impl.ServiceImpl;
import com.tianji.aigc.ChatMemory.RedisChatMemory;
import com.tianji.aigc.dto.ChatDTO;
import com.tianji.aigc.entity.ChatSession;
import com.tianji.aigc.enums.ChatEventTypeEnum;
import com.tianji.aigc.mapper.ChatSessionMapper;
import com.tianji.aigc.service.IChatService;
import com.tianji.aigc.service.IChatSessionService;
import com.tianji.aigc.vo.ChatEventVO;
import com.tianji.common.exceptions.BadRequestException;
import com.tianji.common.utils.UserContext;
import org.springframework.ai.chat.client.ChatClient;
import org.springframework.ai.chat.client.advisor.AbstractChatMemoryAdvisor;
import org.springframework.ai.chat.messages.AssistantMessage;
import org.springframework.beans.factory.annotation.Autowired;
import org.springframework.stereotype.Service;
import reactor.core.publisher.Flux;

import java.util.Map;
import java.util.concurrent.ConcurrentHashMap;

/**
 * 对话业务实现类
 *
 * @Author mr.wu
 * @Date 2025-8-17 15:28
 */
@Service
public class ChatServiceImpl extends ServiceImpl<ChatSessionMapper, ChatSession> implements IChatService {
    @Autowired
    private ChatClient chatClient;

    @Autowired
    private IChatSessionService chatSessionService;

    //定义可以控制连续输出标识值的Map，Map的KEY是sessionId，Vlue就是布尔值，true表示可以继续输出，false表示中断输出.这里采用ConcurrentHashMap是确保线程安全，如果分布式环境可考虑用redis
    private static Map<String,Boolean> SESSION_MAP = new ConcurrentHashMap<String,Boolean>();

    /**
     * 流式对话
     *
     * @param dto
     * @return
     */
    @Override
    public Flux<ChatEventVO> chat(ChatDTO dto) {
        //判断是会话ID是否存在
        ChatSession chatSession = chatSessionService.lambdaQuery().eq(ChatSession::getSessionId, dto.getSessionId()).one();
        if(chatSession==null) {
            throw new BadRequestException("会话不存在");
        }

        //拼接对话ID： 用户ID+会话ID
        String conversationId = UserContext.getUser() + ":" + dto.getSessionId();

        StringBuffer content = new StringBuffer();//线程安全的字符串拼接类，用于拼接流式输出的每个字符串

        return chatClient.prompt()
                .user(dto.getQuestion())
                .advisors(x->x.param(AbstractChatMemoryAdvisor.CHAT_MEMORY_CONVERSATION_ID_KEY,conversationId))
                .stream()
                .content()
                .doFirst(() -> SESSION_MAP.put(dto.getSessionId(),true))//标记开始输出
                .doOnComplete(() -> SESSION_MAP.remove(dto.getSessionId()))//标记结束输出
                .doOnError(throwable -> SESSION_MAP.remove(dto.getSessionId()))//输出异常,清除标记
                .doOnCancel(()-> this.saveRedis(conversationId, content.toString()))//保存停止生成时已有的输出
                .takeWhile(x -> SESSION_MAP.getOrDefault(dto.getSessionId(),false))//判断是否继续输出，false停止输出，true继续输出
                .map(x -> {
                    content.append(x);
                    return ChatEventVO.builder().eventData(content).eventType(ChatEventTypeEnum.DATA.getValue()).build();}
                )
                .concatWith(Flux.just(ChatEventVO.builder().eventType(ChatEventTypeEnum.STOP.getValue()).build()));
    }


    @Autowired
    private RedisChatMemory redisChatMemory;

    private void saveRedis(String conversationId, String content) {
        redisChatMemory.add(conversationId,new AssistantMessage(content));
    }

    /**
     * 停止生成
     * @param sessionId
     */
    @Override
    public void stop(String sessionId) {
        SESSION_MAP.put(sessionId,false);
    }
}
